"""
CN-CLIP 文物数据集基线测试主脚本
"""

import os
import sys
import logging
import argparse
from pathlib import Path
import json
import torch

# 添加项目根目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))

from codebase.cnclip_raw.model_loader import CNClipModelLoader
from codebase.cnclip_raw.dataset_adapter import WenwuRawDataset, create_raw_dataloader
from codebase.cnclip_raw.evaluator import CNClipEvaluator

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def run_baseline_test(args):
    """运行基线测试"""
    logger.info("=" * 80)
    logger.info("CN-CLIP 文物数据集基线测试开始")
    logger.info("=" * 80)
    
    # 1. 加载模型
    logger.info(f"Loading model: {args.model_name}")
    model_loader = CNClipModelLoader(device=args.device)
    
    if not model_loader.load_model(args.model_name, download_root=args.download_root):
        logger.error("Failed to load model")
        return False
    
    logger.info("Model loaded successfully")
    logger.info(f"Model info: {model_loader.get_model_info()}")
    
    # 2. 准备数据集
    logger.info(f"Preparing dataset: {args.start_p:.2f} to {args.end_p:.2f}")
    
    try:
        dataset = WenwuRawDataset(
            start_p=args.start_p,
            end_p=args.end_p,
            preprocess=model_loader.preprocess
        )
        
        logger.info(f"Dataset statistics: {dataset.get_statistics()}")
        
        dataloader = create_raw_dataloader(
            dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers
        )
        
    except Exception as e:
        logger.error(f"Failed to load dataset: {e}")
        return False
    
    # 3. 运行评估
    logger.info("Starting evaluation...")
    evaluator = CNClipEvaluator(model_loader, device=args.device)
    
    try:
        results = evaluator.evaluate_retrieval(
            dataloader,
            k_values=args.k_values,
            max_samples=args.max_samples,
            save_features=args.save_features
        )
        
    except Exception as e:
        logger.error(f"Evaluation failed: {e}")
        return False
    
    # 4. 生成报告
    report = evaluator.generate_report(results, save_path=args.report_path)
    print("\n" + report)
    
    # 5. 保存详细结果 
    if args.results_path:
        # 移除可能很大的特征数据以节省空间
        save_results = results.copy()
        if "features" in save_results:
            del save_results["features"]
        
        with open(args.results_path, 'w', encoding='utf-8') as f:
            json.dump(save_results, f, indent=2, ensure_ascii=False, default=str)
        logger.info(f"Detailed results saved to {args.results_path}")
    
    logger.info("Baseline test completed successfully!")
    return True


def main():
    parser = argparse.ArgumentParser(description="CN-CLIP文物数据集基线测试")
    
    # 模型参数
    parser.add_argument("--model_name", type=str, default="ViT-B-16",
                       help="CN-CLIP模型名称")
    parser.add_argument("--download_root", type=str, default="./pretrained_weights",
                       help="预训练模型下载路径")
    parser.add_argument("--device", type=str, default="auto",
                       help="设备 (auto/cpu/cuda)")
    
    # 数据集参数
    parser.add_argument("--start_p", type=float, default=0.0,
                       help="数据集起始比例")
    parser.add_argument("--end_p", type=float, default=0.1,
                       help="数据集结束比例")
    parser.add_argument("--batch_size", type=int, default=32,
                       help="批大小")
    parser.add_argument("--num_workers", type=int, default=2,
                       help="数据加载线程数")
    
    # 评估参数
    parser.add_argument("--k_values", type=int, nargs="+", default=[1, 5, 10],
                       help="Recall@K的K值列表")
    parser.add_argument("--max_samples", type=int, default=None,
                       help="最大样本数")
    parser.add_argument("--save_features", action="store_true",
                       help="是否保存特征向量")
    
    # 输出参数
    parser.add_argument("--report_path", type=str, default="research/docs/CN-CLIP基线测试报告.md",
                       help="报告保存路径")
    parser.add_argument("--results_path", type=str, default="research/docs/CN-CLIP基线测试结果.json",
                       help="详细结果保存路径")
    
    args = parser.parse_args()
    
    # 自动选择设备
    if args.device == "auto":
        args.device = "cuda" if torch.cuda.is_available() else "cpu"
    
    logger.info(f"Arguments: {vars(args)}")
    
    # 创建输出目录
    os.makedirs(os.path.dirname(args.report_path) if os.path.dirname(args.report_path) else ".", exist_ok=True)
    os.makedirs(os.path.dirname(args.results_path) if os.path.dirname(args.results_path) else ".", exist_ok=True)
    
    # 运行测试
    success = run_baseline_test(args)
    
    if success:
        logger.info("All tests completed successfully!")
        sys.exit(0)
    else:
        logger.error("Tests failed!")
        sys.exit(1)


if __name__ == "__main__":
    main()